"""TestCase class with benefits."""
import contextlib
import sys
import typing
import unittest

_T = typing.TypeVar("_T")


class TestCaseWithContext(unittest.TestCase):
    """TestCase with some backports."""

    if sys.version_info < (3, 11):

        def enterContext(self, cm: contextlib.AbstractContextManager[_T]) -> _T:
            """Enters the supplied context manager."""
            result = cm.__enter__()  # pylint: disable=unnecessary-dunder-call
            self.addCleanup(cm.__exit__, None, None, None)
            return result

        @classmethod
        def enterClassContext(cls, cm: contextlib.AbstractContextManager[_T]) -> _T:
            """Enters the supplied class-wide context manager."""
            result = cm.__enter__()  # pylint: disable=unnecessary-dunder-call
            cls.addClassCleanup(cm.__exit__, None, None, None)
            return result
